#include <webots/Camera.hpp>
#include <webots/Display.hpp>
#include <webots/Keyboard.hpp>
#include <webots/Motor.hpp>
#include <webots/Robot.hpp>

#include <opencv2/core/core.hpp>
#include <opencv2/imgproc/imgproc.hpp>

#include <memory>

#include <depth_model.h>

using namespace webots;

int main([[maybe_unused]] int argc, [[maybe_unused]] char* argv[]) {
  // Instantiate objects
  auto robot = std::make_shared<Robot>();
  auto keyboard = std::make_shared<Keyboard>();

  // Configure camera
  auto* camera = robot->getCamera("camera_top");
  int sampling_period = 32;
  camera->enable(sampling_period);
  // get camera model parameters
  auto cam_width = camera->getWidth();
  auto cam_height = camera->getHeight();

  // configure motors
  std::vector<Motor*> wheels;
  std::vector<std::string> wheels_names = {"wheel_motor_1", "wheel_motor_2", "wheel_motor_3", "wheel_motor_4"};
  for (auto& wheel_name : wheels_names) {
    auto* motor = robot->getMotor(wheel_name);
    wheels.push_back(motor);
    motor->setPosition(std::numeric_limits<double>::infinity());
    motor->setVelocity(0.0);
  }

  // configure display
  auto depth_view = robot->getDisplay("depth_view");

  // milliseconds
  int time_step = 64;  //(int)robot->getBasicTimeStep();

  mono_depth::DepthAnythingModel depth_model("depth_anything_vits14.engine");

  keyboard->enable(time_step);
  const double MAX_SPEED = 2.0;
  double left_speed = 0.0;
  double right_speed = 0.0;
  while (robot->step(time_step) != -1) {
    // capture camera image
    auto* image_data = camera->getImage();  // BGRA (32 bits)
    cv::Mat img = cv::Mat(cv::Size(cam_width, cam_height), CV_8UC4);
    auto data_len = img.total() * img.elemSize();
    memcpy(img.data, image_data, data_len);

    // keyboard control
    auto key = keyboard->getKey();
    switch (key) {
      case Keyboard::UP:
        left_speed = MAX_SPEED;
        right_speed = MAX_SPEED;
        break;
      case Keyboard::DOWN:
        left_speed = -MAX_SPEED;
        right_speed = -MAX_SPEED;
        break;
      case Keyboard::LEFT:
        left_speed = -MAX_SPEED;
        right_speed = MAX_SPEED;
        break;
      case Keyboard::RIGHT:
        left_speed = MAX_SPEED;
        right_speed = -MAX_SPEED;
        break;
      default:
        left_speed = 0.0;
        right_speed = 0.0;
    }
    if (key == 'I') {
    }

    // update motors
    wheels[0]->setVelocity(left_speed);
    wheels[1]->setVelocity(right_speed);
    wheels[2]->setVelocity(left_speed);
    wheels[3]->setVelocity(right_speed);

    // process and show
    {
      cv::cvtColor(img, img, cv::COLOR_BGRA2RGB);
      cv::Mat depth_img = depth_model.process(img);

      auto* ir = depth_view->imageNew(depth_img.cols, depth_img.rows, depth_img.data, Display::RGB);
      depth_view->imagePaste(ir, 0, 0, false);
      depth_view->imageDelete(ir);
    }
  }
  return 0;
}